import os
import json
import requests
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from docx import Document
import re


@dataclass
class DocumentInfo:
    """文档信息类"""
    file_path: str
    file_name: str
    content: str


class AIDocumentSearcher:
    """基于AI API的文档搜索器"""

    def __init__(self, api_key: str, model_name: str = "gpt-4o-mini"):
        """
        初始化搜索器

        Args:
            api_key: GeeKAI API密钥
            model_name: 使用的模型名称
        """
        self.documents: List[DocumentInfo] = []
        self.api_key = api_key
        self.model_name = model_name
        self.api_base_url = "https://geekai.co/api/v1/chat/completions"
        self.last_query = ""
        self.found_relevant_docs = []

    def extract_docx_content(self, file_path: str) -> str:
        """
        提取docx文件内容

        Args:
            file_path: 文件路径

        Returns:
            str: 文件内容
        """
        try:
            doc = Document(file_path)
            content = []

            # 提取段落文本
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    content.append(paragraph.text.strip())

            # 提取表格文本
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        content.append(" | ".join(row_text))

            return "\n".join(content)

        except Exception as e:
            print(f"❌ 读取文件 {file_path} 时出错: {str(e)}")
            return ""

    def scan_directory(self, directory_path: str) -> None:
        """
        扫描目录中的所有docx文件

        Args:
            directory_path: 目录路径
        """
        print(f"🔍 开始扫描目录: {directory_path}")

        docx_files = []

        # 收集所有docx文件
        for root, dirs, files in os.walk(directory_path):
            for file in files:
                if file.lower().endswith('.docx') and not file.startswith('~$'):
                    docx_files.append(os.path.join(root, file))

        print(f"📁 找到 {len(docx_files)} 个 .docx 文件")

        # 处理每个文件
        for i, file_path in enumerate(docx_files, 1):
            print(f"📄 处理文件 ({i}/{len(docx_files)}): {os.path.basename(file_path)}")

            content = self.extract_docx_content(file_path)
            if content:
                doc_info = DocumentInfo(
                    file_path=file_path,
                    file_name=os.path.basename(file_path),
                    content=content
                )
                self.documents.append(doc_info)

        print(f"✅ 扫描完成，共处理了 {len(self.documents)} 个文档")

    def call_ai_api(self, messages: List[Dict], max_tokens: int = 1000) -> str:
        """
        调用GeeKAI API

        Args:
            messages: 消息列表
            max_tokens: 最大token数

        Returns:
            str: AI响应
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        data = {
            "model": self.model_name,
            "messages": messages,
            "max_tokens": max_tokens
        }

        try:
            response = requests.post(
                self.api_base_url,
                headers=headers,
                json=data,
                timeout=60
            )

            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content']
            else:
                return f"❌ API调用失败 (状态码: {response.status_code}): {response.text}"

        except requests.exceptions.Timeout:
            return "❌ API调用超时"
        except Exception as e:
            return f"❌ API调用出错: {str(e)}"

    def send_documents_to_ai(self) -> str:
        """
        将所有文档信息发送给AI

        Returns:
            str: AI确认消息
        """
        if not self.documents:
            return "❌ 没有文档可处理"

        print("🤖 正在将文档信息发送给AI...")

        # 构建文档信息摘要
        doc_summaries = []
        for i, doc_info in enumerate(self.documents, 1):
            # 限制每个文档的内容长度用于初始化
            content_preview = doc_info.content[:1000] + "..." if len(doc_info.content) > 1000 else doc_info.content

            doc_summary = f"""
文档{i}:
文件名: {doc_info.file_name}
文件路径: {doc_info.file_path}
内容: {content_preview}
---
"""
            doc_summaries.append(doc_summary)

        all_docs_content = "\n".join(doc_summaries)

        messages = [
            {
                "role": "system",
                "content": """你是一个专业的文档管理助手。我将向你提供一批文档的信息，包括文件名、路径和内容。
请记住这些文档的内容，当用户提问时，你需要：
1. 准确找到相关的文档
2. 指出具体的文件名和路径
3. 引用相关的内容片段
4. 分析相关性原因

请确认你已经接收并理解了所有文档信息。"""
            },
            {
                "role": "user",
                "content": f"以下是所有文档信息：\n\n{all_docs_content}\n\n请确认你已经接收并分析了这些文档。"
            }
        ]

        response = self.call_ai_api(messages, max_tokens=500)
        print("✅ 文档信息已发送给AI")
        return response

    def analyze_document_relevance(self, doc_info: DocumentInfo, query: str) -> str:
        """
        分析文档与查询的相关性

        Args:
            doc_info: 文档信息
            query: 查询内容

        Returns:
            str: 相关性分析结果
        """
        # 限制文档内容长度以避免token超限
        content = doc_info.content[:2000] + "..." if len(doc_info.content) > 2000 else doc_info.content

        messages = [
            {
                "role": "system",
                "content": """你是一个文档分析专家。请分析给定文档与用户查询的相关性，并具体说明：
1. 文档中哪些部分与查询相关
2. 为什么相关（相关性原因）
3. 相关的具体内容片段
4. 相关性的强弱程度

请用简洁明了的语言回答。"""
            },
            {
                "role": "user",
                "content": f"""
查询: {query}

文档内容:
{content}

请分析这个文档与查询"{query}"的相关性，具体说明相关的部分和原因。
"""
            }
        ]

        return self.call_ai_api(messages, max_tokens=400)

    def analyze_missing_content(self, query: str, relevant_docs: List[str]) -> str:
        """
        分析查询中缺失的内容并提供补充建议

        Args:
            query: 原始查询
            relevant_docs: 找到的相关文档列表

        Returns:
            str: 缺失内容分析和补充建议
        """
        print("🔍 分析缺失内容和补充建议...")

        # 构建已找到的相关文档信息
        found_docs_info = []
        for doc_name in relevant_docs:
            for doc in self.documents:
                if doc.file_name == doc_name:
                    content_preview = doc.content[:800] + "..." if len(doc.content) > 800 else doc.content
                    found_docs_info.append(f"文档: {doc.file_name}\n内容概要: {content_preview}")
                    break

        # 构建所有文档的概要信息
        all_docs_summary = []
        for doc in self.documents:
            doc_summary = f"- {doc.file_name}: {doc.content[:200]}..." if len(
                doc.content) > 200 else f"- {doc.file_name}: {doc.content}"
            all_docs_summary.append(doc_summary)

        found_content = "\n".join(found_docs_info) if found_docs_info else "无相关文档"
        all_docs_overview = "\n".join(all_docs_summary)

        messages = [
            {
                "role": "system",
                "content": """你是一个文档管理和内容规划专家。请分析用户查询中哪些内容在现有文档中没有找到，并提供具体的补充建议。

你的分析应该包括：
1. 查询内容的拆解分析
2. 哪些部分在现有文档中已经覆盖
3. 哪些部分缺失或不完整
4. 具体的补充建议：
   - 可以在哪个现有文档中添加什么内容
   - 或者需要新建什么主题的文档
   - 具体应该包含哪些内容要点

请提供实用且具体的建议。"""
            },
            {
                "role": "user",
                "content": f"""
用户查询: {query}

已找到的相关文档内容:
{found_content}

现有所有文档概览:
{all_docs_overview}

请分析：
1. 用户查询"{query}"中哪些内容在现有文档中没有完全覆盖？
2. 针对缺失的内容，应该如何补充？可以在现有哪个文档中添加什么内容，或者需要新建什么文档？
3. 请提供具体的内容补充建议和文档规划建议。
"""
            }
        ]

        return self.call_ai_api(messages, max_tokens=800)

    def extract_file_names_from_ai_response(self, ai_response: str) -> List[str]:
        """
        从AI响应中提取文件名

        Args:
            ai_response: AI响应内容

        Returns:
            List[str]: 文件名列表
        """
        file_names = []

        # 查找AI响应中提到的文档
        for doc in self.documents:
            if doc.file_name in ai_response:
                file_names.append(doc.file_name)

        return file_names

    def search_with_ai_and_show_content(self, query: str) -> None:
        """
        使用AI进行智能搜索并显示相关文档内容和相关性分析

        Args:
            query: 查询内容
        """
        self.last_query = query
        print(f"🤖 AI搜索中...")

        # 第一步：让AI找到相关文档
        doc_list = []
        for i, doc_info in enumerate(self.documents, 1):
            doc_list.append(f"文档{i}: {doc_info.file_name} (路径: {doc_info.file_path})")

        doc_context = "\n".join(doc_list)

        messages = [
            {
                "role": "system",
                "content": f"""你是一个文档搜索助手。你已经接收了以下文档信息：

{doc_context}

当用户提问时，请找出与查询相关的所有文档，并在回答中明确列出所有相关文档的文件名。"""
            },
            {
                "role": "user",
                "content": f"查询：{query}\n\n请找出与此查询相关的所有文档，并明确说明每个文档的文件名。"
            }
        ]

        ai_result = self.call_ai_api(messages, max_tokens=600)
        print(f"\n🔍 AI搜索结果:\n{ai_result}")

        # 第二步：提取相关文档的文件名
        relevant_file_names = self.extract_file_names_from_ai_response(ai_result)
        self.found_relevant_docs = relevant_file_names

        if not relevant_file_names:
            print("\n❌ 未找到相关文档")
            # 即使没找到相关文档，也要分析缺失内容
            self.analyze_and_show_missing_content(query, [])
            return

        # 第三步：分析每个相关文档并显示内容
        print(f"\n📚 找到 {len(relevant_file_names)} 个相关文档:")
        print("=" * 100)

        for i, file_name in enumerate(relevant_file_names, 1):
            doc = None
            for document in self.documents:
                if document.file_name == file_name:
                    doc = document
                    break

            if doc:
                print(f"\n📄 文档 {i}: {doc.file_name}")
                print(f"📁 路径: {doc.file_path}")
                print(f"📊 内容长度: {len(doc.content)} 字符")
                print("-" * 50)

                # 显示原文内容（限制100字符）
                content_preview = doc.content[:100]
                if len(doc.content) > 100:
                    content_preview += "..."

                print(f"📖 原文内容: {content_preview}")
                print("-" * 50)

                # 分析相关性
                print("🔍 相关性分析:")
                relevance_analysis = self.analyze_document_relevance(doc, query)
                print(relevance_analysis)
                print("=" * 100)

                # 如果有多个文档，询问是否继续
                if i < len(relevant_file_names):
                    continue_show = input(f"\n⏩ 是否继续查看下一个相关文档？(y/n，直接回车继续): ").strip().lower()
                    if continue_show == 'n':
                        print("📋 剩余相关文档列表:")
                        for j in range(i, len(relevant_file_names)):
                            remaining_doc = None
                            for document in self.documents:
                                if document.file_name == relevant_file_names[j]:
                                    remaining_doc = document
                                    break
                            if remaining_doc:
                                print(f"   {j + 1}. {remaining_doc.file_name} ({remaining_doc.file_path})")
                        break

        print("\n✅ 所有相关文档分析完毕")

        # 第四步：分析缺失内容并提供补充建议
        self.analyze_and_show_missing_content(query, relevant_file_names)

    def analyze_and_show_missing_content(self, query: str, relevant_docs: List[str]) -> None:
        """
        分析并显示缺失内容和补充建议

        Args:
            query: 查询内容
            relevant_docs: 相关文档列表
        """
        print("\n" + "🔍" * 40)
        print("📋 缺失内容分析与补充建议")
        print("🔍" * 40)

        missing_analysis = self.analyze_missing_content(query, relevant_docs)
        print(missing_analysis)

        print("\n" + "🔍" * 40)
        print("✅ 补充建议分析完毕")
        print("🔍" * 40)

    def interactive_search(self) -> None:
        """
        交互式搜索模式
        """
        # 首先发送文档信息给AI
        ai_confirmation = self.send_documents_to_ai()
        print(f"🤖 AI响应: {ai_confirmation}")

        print("\n" + "=" * 80)
        print("📚 AI文档智能搜索系统已准备就绪！")
        print(f"📊 已加载 {len(self.documents)} 个文档")
        print("\n💡 使用说明:")
        print("   - 直接输入要搜索的内容")
        print("   - AI会智能分析并找到相关文档")
        print("   - 显示原文内容（前100字符）和详细相关性分析")
        print("   - 🆕 自动分析缺失内容并提供补充建议")
        print("   - 支持模糊搜索和语义理解")
        print("\n🔧 命令:")
        print("   - 直接输入查询内容")
        print("   - 输入 'list' 查看所有文档")
        print("   - 输入 'help' 查看帮助")
        print("   - 输入 'missing' 重新分析上次查询的缺失内容")
        print("   - 输入 'quit' 退出")
        print("=" * 80)

        while True:
            try:
                user_input = input("\n🔍 请输入搜索内容: ").strip()

                if not user_input:
                    continue

                if user_input.lower() == 'quit':
                    print("👋 再见！")
                    break

                if user_input.lower() == 'help':
                    print("""
📖 详细帮助信息:

🔸 搜索示例:
   - "床头明月光" - AI会找到相关诗词文档，显示内容和相关性分析
   - "项目管理流程" - 查找项目管理相关文档，分析为什么相关
   - "财务报表分析" - 查找财务相关文档，说明相关的具体部分

🔸 系统特点:
   - 智能相关性分析：详细说明文档为什么与查询相关
   - 内容预览：显示文档前100字符内容
   - 精确定位：指出具体相关的内容片段
   - 🆕 缺失内容分析：分析查询中哪些内容在现有文档中没有
   - 🆕 补充建议：提供具体的文档补充和扩展建议

🔸 显示内容:
   - 文档基本信息（文件名、路径、长度）
   - 原文内容前100字符（超出用...显示）
   - 详细的相关性分析和原因说明
   - 具体相关的内容片段
   - 🆕 缺失内容分析和补充建议

🔸 补充建议功能:
   - 分析查询内容的各个组成部分
   - 识别现有文档中未覆盖的内容
   - 建议在现有文档中添加什么内容
   - 建议新建什么主题的文档
   - 提供具体的内容要点建议

🔸 其他命令:
   - 'list' - 显示所有已加载的文档
   - 'missing' - 重新分析上次查询的缺失内容
   - 'quit' - 退出程序
                    """)
                    continue

                if user_input.lower() == 'list':
                    print(f"\n📋 已加载的文档列表 (共{len(self.documents)}个):")
                    for i, doc in enumerate(self.documents, 1):
                        print(f"   {i}. {doc.file_name}")
                        print(f"      路径: {doc.file_path}")
                        print(f"      内容长度: {len(doc.content)} 字符")
                    continue

                if user_input.lower() == 'missing':
                    if self.last_query:
                        print(f"🔍 重新分析查询 '{self.last_query}' 的缺失内容...")
                        self.analyze_and_show_missing_content(self.last_query, self.found_relevant_docs)
                    else:
                        print("❌ 请先进行一次搜索，然后再使用此功能")
                    continue

                # 普通搜索
                self.search_with_ai_and_show_content(user_input)

            except KeyboardInterrupt:
                print("\n👋 程序已退出")
                break
            except Exception as e:
                print(f"❌ 搜索时出错: {str(e)}")


def main():
    """主程序"""
    print("🚀 启动AI文档智能搜索系统")
    print("=" * 50)

    # 获取API密钥
    api_key = input("请输入GeeKAI API密钥: ").strip()
    if not api_key:
        print("❌ API密钥不能为空")
        return

    # 获取模型名称
    model = input("请输入模型名称 (直接回车使用 gpt-4o-mini): ").strip()
    if not model:
        model = "gpt-4o-mini"

    # 获取要扫描的目录
    directory_path = input("请输入要扫描的文件夹路径: ").strip()

    if not os.path.exists(directory_path):
        print("❌ 路径不存在，请检查后重试")
        return

    # 初始化搜索器
    searcher = AIDocumentSearcher(api_key=api_key, model_name=model)

    # 扫描文档
    searcher.scan_directory(directory_path)

    if not searcher.documents:
        print("❌ 未找到任何 .docx 文件")
        return

    # 开始交互式搜索
    searcher.interactive_search()


if __name__ == "__main__":
    main()
